#include <dirent.h>
#include <sys/types.h>
#include <unistd.h>

#include <algorithm>
#include <boost/filesystem.hpp>
#include <chrono>
#include <ctime>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <nlohmann_json/json.hpp>  // 引入 nlohmann::json
#include <sstream>
#include <string>
#include <thread>
#include <vector>

namespace fs = boost::filesystem;  // 使用 boost::filesystem

using json = nlohmann::json;  // 使用简化的 json 名称

// Function to read /proc/meminfo
void readMemInfo(double &totalMem, double &usedMem, double &freeMem) {
  std::ifstream meminfo("/proc/meminfo");
  std::string line;
  double total = 0, free = 0, buffers = 0, cached = 0;

  while (std::getline(meminfo, line)) {
    std::istringstream iss(line);
    std::string key;
    double value;
    std::string unit;
    iss >> key >> value >> unit;

    if (key == "MemTotal:") {
      total = value;
    } else if (key == "MemFree:") {
      free = value;
    } else if (key == "Buffers:") {
      buffers = value;
    } else if (key == "Cached:") {
      cached = value;
    }
  }

  // Convert to MB
  totalMem = total / 1024.0;
  freeMem = free / 1024.0;
  usedMem = (total - free - buffers - cached) / 1024.0;
}

// Function to read CPU usage statistics from /proc/stat
void readCpuStats(std::map<int, std::vector<double>> &cpuStats) {
  std::ifstream stat("/proc/stat");
  std::string line;

  while (std::getline(stat, line)) {
    std::istringstream iss(line);
    std::string cpu;
    double user, nice, system, idle, iowait, irq, softirq, steal;
    if (iss >> cpu && cpu.rfind("cpu", 0) == 0) {
      if (cpu == "cpu") {
        continue;  // Skip aggregate line
      }
      iss >> user >> nice >> system >> idle >> iowait >> irq >> softirq >>
          steal;
      cpuStats[std::stoi(cpu.substr(3))] = {user,   nice, system,  idle,
                                            iowait, irq,  softirq, steal};
    }
  }
}

// Function to calculate CPU load percentages
void calculateCpuLoad(const std::map<int, std::vector<double>> &currentCpuStats,
                      std::map<int, double> &cpuLoad, double &totalCpuLoad) {
  static std::map<int, std::vector<double>> previousCpuStats;
  static std::map<int,
                  std::chrono::time_point<std::chrono::high_resolution_clock>>
      previousTimestamps;

  auto currentTime = std::chrono::high_resolution_clock::now();
  double totalIdle = 0, totalNonIdle = 0, totalCpuTime = 0;

  for (const auto &entry : currentCpuStats) {
    int core = entry.first;
    const auto &current = entry.second;

    if (previousCpuStats.find(core) != previousCpuStats.end()) {
      const auto &previous = previousCpuStats[core];
      double totalCurrent = current[0] + current[1] + current[2] + current[3] +
                            current[4] + current[5] + current[6] + current[7];
      double totalPrevious = previous[0] + previous[1] + previous[2] +
                             previous[3] + previous[4] + previous[5] +
                             previous[6] + previous[7];
      double idleCurrent = current[3];
      double idlePrevious = previous[3];

      auto timeDiff = std::chrono::duration_cast<std::chrono::seconds>(
                          currentTime - previousTimestamps[core])
                          .count();
      if (timeDiff > 0) {
        double totalDelta = totalCurrent - totalPrevious;
        double idleDelta = idleCurrent - idlePrevious;
        double cpuLoadPercentage =
            (totalDelta > 0) ? ((totalDelta - idleDelta) / totalDelta) * 100.0
                             : 0.0;

        cpuLoad[core] = cpuLoadPercentage;
        totalCpuTime += totalDelta;
        totalIdle += idleDelta;
        totalNonIdle += (totalDelta - idleDelta);
      } else {
        cpuLoad[core] = 0.0;
      }
    } else {
      cpuLoad[core] = 0.0;
    }

    previousCpuStats[core] = current;
    previousTimestamps[core] = currentTime;
  }

  if (totalCpuTime > 0) {
    totalCpuLoad = (totalCpuTime > 0)
                       ? ((totalCpuTime - totalIdle) / totalCpuTime) * 100.0
                       : 0.0;
  } else {
    totalCpuLoad = 0.0;
  }
}

// Function to read process memory usage from /proc/[pid]/status
double getProcessMemoryUsage(int pid) {
  std::ifstream statusFile("/proc/" + std::to_string(pid) + "/status");
  std::string line;
  double vmRSS = 0;
  while (std::getline(statusFile, line)) {
    if (line.find("VmRSS:") != std::string::npos) {
      std::istringstream iss(line);
      std::string key;
      iss >> key >> vmRSS;
      return vmRSS / 1024.0;  // Convert to MB
    }
  }
  return 0.0;
}

// Function to calculate process CPU usage
void getProcessCpuUsage(int pid, double &cpuUsage) {
  static std::map<int, std::pair<double, double>> previousTimes;
  static std::map<int,
                  std::chrono::time_point<std::chrono::high_resolution_clock>>
      previousTimestamps;

  std::ifstream statFile("/proc/" + std::to_string(pid) + "/stat");
  std::string line;
  if (std::getline(statFile, line)) {
    std::istringstream iss(line);
    std::string token;
    for (int i = 0; i < 13; ++i) {
      iss >> token;
    }
    double utime, stime;
    iss >> utime >> stime;
    auto currentTime = std::chrono::high_resolution_clock::now();

    if (previousTimes.find(pid) != previousTimes.end()) {
      double previousUtime = previousTimes[pid].first;
      double previousStime = previousTimes[pid].second;
      auto timeDiff = std::chrono::duration_cast<std::chrono::seconds>(
                          currentTime - previousTimestamps[pid])
                          .count();

      if (timeDiff > 0) {
        double utimeDiff = utime - previousUtime;
        double stimeDiff = stime - previousStime;
        cpuUsage =
            ((utimeDiff + stimeDiff) / sysconf(_SC_CLK_TCK)) * 100.0 / timeDiff;
        unsigned int numCores = std::thread::hardware_concurrency();
        cpuUsage = std::round(cpuUsage / numCores * 10) / 10;
      } else {
        cpuUsage = 0.0;
      }
    } else {
      cpuUsage = 0.0;
    }

    previousTimes[pid] = {utime, stime};
    previousTimestamps[pid] = currentTime;
  } else {
    cpuUsage = 0.0;
  }
}

// Function to read JSON file and parse process names, max files, and output
// file
void readProcessNamesAndMaxFiles(const std::string &filename,
                                 std::vector<std::string> &processNames,
                                 int &maxFiles, std::string &outputFile) {
  std::ifstream file(filename);
  json j;
  file >> j;

  for (const auto &name : j["processes"]) {
    processNames.push_back(name.get<std::string>());
  }

  maxFiles = j["maxFiles"].get<int>();
  outputFile = j["outputFile"].get<std::string>();
}

// Function to find all PIDs matching the process names
std::vector<std::pair<int, std::string>> findPids(
    const std::vector<std::string> &processNames) {
  std::vector<std::pair<int, std::string>> pids;
  DIR *dir = opendir("/proc");
  if (dir == nullptr) {
    std::cerr << "Failed to open /proc directory" << std::endl;
    return pids;
  }

  struct dirent *entry;
  while ((entry = readdir(dir)) != nullptr) {
    if (entry->d_type == DT_DIR) {
      std::string dirName(entry->d_name);
      if (isdigit(dirName[0])) {
        int pid = std::stoi(dirName);
        std::ifstream cmdlineFile("/proc/" + dirName + "/cmdline");
        std::string cmdline;
        std::getline(cmdlineFile, cmdline);

        for (const auto &processName : processNames) {
          if (cmdline.find(processName) != std::string::npos) {
            pids.emplace_back(pid, cmdline);
          }
        }
      }
    }
  }
  closedir(dir);
  return pids;
}

// Function to generate timestamped file name
std::string generateFileName(const std::string &basePath) {
  std::string fileName;
  do {
    auto now = std::chrono::system_clock::now();
    std::time_t now_c = std::chrono::system_clock::to_time_t(now);
    std::tm tm = *std::localtime(&now_c);

    std::ostringstream oss;
    oss << basePath << "/resource_statistics_"
        << std::put_time(&tm, "%Y%m%d_%H%M%S") << ".txt";
    fileName = oss.str();

    if (boost::filesystem::exists(fileName)) {
      std::this_thread::sleep_for(std::chrono::seconds(1));
    }
  } while (boost::filesystem::exists(fileName));

  return fileName;
}

// Function to manage files (create a new file if necessary)
void manageFiles(const std::string &basePath, int maxFiles,
                 std::string &currentFile) {
  if (!boost::filesystem::exists(basePath)) {
    if (!boost::filesystem::create_directories(basePath)) {
      std::cerr << "Failed to create directory: " << basePath << std::endl;
      return;
    }
  }

  fs::directory_iterator end_itr;
  std::vector<fs::path> files;

  for (fs::directory_iterator itr(basePath); itr != end_itr; ++itr) {
    if (fs::is_regular_file(itr->path()) && itr->path().extension() == ".txt") {
      files.push_back(itr->path());
    }
  }

  if (files.size() > maxFiles) {
    std::sort(files.begin(), files.end(),
              [](const fs::path &a, const fs::path &b) {
                return fs::last_write_time(a) < fs::last_write_time(b);
              });

    for (std::size_t i = 0; i < files.size() - maxFiles; ++i) {
      fs::remove(files[i]);
    }
  }

  if (currentFile.empty()) {
    currentFile = generateFileName(basePath);
  }
}

// Function to save statistics to the current file
void saveStatistics(const std::string &currentFile, double totalCpuLoad,
                    const std::map<int, double> &cpuLoad, double totalMem,
                    double usedMem, double freeMem,
                    const std::vector<std::pair<int, std::string>> &processes) {
  // Open the file to append data
  std::ofstream outfile(currentFile, std::ios::app);
  if (!outfile.is_open()) {
    std::cerr << "Failed to open file: " << currentFile << std::endl;
    return;
  }

  auto now = std::chrono::system_clock::now();
  std::time_t now_c = std::chrono::system_clock::to_time_t(now);
  std::tm tm = *std::localtime(&now_c);

  outfile << "Timestamp: " << std::put_time(&tm, "%Y-%m-%d %H:%M:%S") << "\n";
  outfile << "MEM[Total, Used, Free]: " << std::fixed << std::setprecision(2)
          << totalMem << " MB, " << usedMem << " MB, " << freeMem << " MB\n";
  outfile << "CPU[Total]: " << std::fixed << std::setprecision(2)
          << totalCpuLoad << " %CPU\n";

  for (const auto &process : processes) {
    double memUsage = getProcessMemoryUsage(process.first);
    double cpuUsage;
    getProcessCpuUsage(process.first, cpuUsage);
    std::string process_name = process.second;
    // std::replace(process_name.begin(), process_name.end(), '^', ' ');
    // std::replace(process_name.begin(), process_name.end(), '@', ' ');
    std::replace(process_name.begin(), process_name.end(), '\0', ' ');
    outfile << process_name << "[Pid, cpu, mem]: " << process.first << ", "
            << std::fixed << std::setprecision(2) << cpuUsage << " %, "
            << memUsage << " MB\n";
  }

  outfile << "\n";  // Add empty line between entries
}

int main(int argc, char *argv[]) {
  // Configuration file path
  std::string configFilePath =
      "config/param.json";  // Specify your JSON file path here
  if (argc == 2) {
    configFilePath.clear();
    configFilePath.assign(argv[1]);
  }
  std::vector<std::string> processNames;
  int maxFiles;
  std::string outputFilePath;
  readProcessNamesAndMaxFiles(configFilePath, processNames, maxFiles,
                              outputFilePath);

  std::string currentFile = "";
  manageFiles(outputFilePath, maxFiles, currentFile);

  while (true) {
    // Read memory info
    double totalMem, usedMem, freeMem;
    readMemInfo(totalMem, usedMem, freeMem);

    // Read CPU stats
    std::map<int, std::vector<double>> cpuStats;
    readCpuStats(cpuStats);

    // Calculate CPU loads
    std::map<int, double> cpuLoad;
    double totalCpuLoad;
    calculateCpuLoad(cpuStats, cpuLoad, totalCpuLoad);

    // Find PIDs
    std::vector<std::pair<int, std::string>> processes = findPids(processNames);

    // Save statistics
    saveStatistics(currentFile, totalCpuLoad, cpuLoad, totalMem, usedMem,
                   freeMem, processes);

    // Sleep for 1 second
    std::this_thread::sleep_for(std::chrono::seconds(1));
  }

  return 0;
}
